#!/bin/bash
# 内存泄漏修复脚本（支持Java应用）
# 适用于Linux系统中检测到内存泄漏(对象未释放)的情况
# 知识图谱关联配置 - 与故障节点ID对应
FIX_CODE_ID="memory-leak-001"  # 对应知识图谱中"内存泄漏（对象未释放）"故障ID


# 配置参数
MEMORY_THRESHOLD=80          # 系统内存使用率阈值(%)
PROCESS_MEM_THRESHOLD=10     # 单个进程内存使用率阈值(%)
CHECK_INTERVAL=30            # 检查间隔时间(秒)
LEAK_CHECK_COUNT=3           # 确认内存泄漏所需的连续检查次数
LOG_FILE="/var/log/memory_leak_fix.log"
WHITELIST_PROCESSES=("init" "systemd" "kernel" "sshd" "bash" "docker" "kubelet")  # 系统关键进程白名单

# Java应用特殊配置
JAVA_PROCESSES=("my-java-service" "spring-boot-app" "tomcat" "myapp.jar")  # Java应用标识
JAVA_HEAP_DUMP_DIR="/var/log/java_heap_dumps"  # Java堆转储文件保存目录

# 服务重启命令映射 - 包含Java应用
SERVICE_RESTART_CMD=(
    "nginx:systemctl restart nginx"
    "mysql:systemctl restart mysql"
    "java:systemctl restart my-java-service"  # 通用Java服务
    "spring-boot-app:/opt/myapp/stop.sh && /opt/myapp/start.sh"  # Spring Boot应用
    "catalina.sh:systemctl restart tomcat"  # Tomcat
    "myapp.jar:/usr/bin/java -jar /opt/myapp/myapp.jar"  # JAR包应用
)

# 初始化日志
log() {
    local timestamp=$(date "+%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] $1" >> $LOG_FILE
    echo "[$timestamp] $1"  # 同时输出到控制台
}

# 创建Java堆转储目录
init_java_heap_dump_dir() {
    if [ ! -d "$JAVA_HEAP_DUMP_DIR" ]; then
        mkdir -p "$JAVA_HEAP_DUMP_DIR"
        chmod 755 "$JAVA_HEAP_DUMP_DIR"
        log "创建Java堆转储目录: $JAVA_HEAP_DUMP_DIR"
    fi
}

# 获取系统内存使用率
get_system_memory_usage() {
    local usage=$(free | awk '/Mem/{printf("%.2f"), $3/$2*100}')
    echo "$usage"
}

# 获取高内存占用的进程列表（增强Java进程识别）
get_high_memory_processes() {
    # 获取内存使用率超过阈值的进程，排除内核进程和白名单进程
    ps -eo pid,user,%mem,rss,vsize,cmd --sort=-%mem | awk -v threshold=$PROCESS_MEM_THRESHOLD '$3 > threshold {print}' | while read -r pid user mem rss vsize cmd; do
        local is_whitelisted=0
        for process in "${WHITELIST_PROCESSES[@]}"; do
            if echo "$cmd" | grep -q "$process"; then
                is_whitelisted=1
                break
            fi
        done
        # 特殊处理Java进程，避免被误判为白名单
        if echo "$cmd" | grep -q "java" && ! echo "$cmd" | grep -q -E "$(IFS=\|; echo "${WHITELIST_PROCESSES[*]}")"; then
            is_whitelisted=0
        fi
        if [ $is_whitelisted -eq 0 ] && [ $pid -ne $$ ]; then  # 排除脚本自身
            echo "$pid $user $mem $rss $vsize $cmd"
        fi
    done
}

# 为Java进程创建堆转储（用于事后分析内存泄漏）
create_java_heap_dump() {
    local pid=$1
    local process_name=$2
    
    # 检查是否为Java进程
    if ps -p $pid -o cmd --no-headers | grep -q "java"; then
        log "为Java进程 $pid ($process_name) 创建堆转储文件"
        
        # 检查jmap命令是否存在
        if command -v jmap &> /dev/null; then
            local dump_file="$JAVA_HEAP_DUMP_DIR/java_heap_dump_${process_name}_${pid}_$(date +%Y%m%d_%H%M%S).hprof"
            jmap -dump:format=b,file=$dump_file $pid
            if [ $? -eq 0 ]; then
                log "Java堆转储文件已保存至: $dump_file"
                return 0
            else
                log "创建Java堆转储文件失败"
                return 1
            fi
        else
            log "未找到jmap命令，无法创建Java堆转储（请安装JDK）"
            return 1
        fi
    fi
    return 0
}

# 检查进程是否存在内存泄漏迹象
check_memory_leak() {
    local pid=$1
    local check_count=$2
    local interval=$3
    
    # 检查进程是否存在
    if ! ps -p $pid > /dev/null; then
        log "进程 $pid 不存在，无法检查内存泄漏"
        return 1
    fi
    
    local process_name=$(ps -p $pid -o cmd --no-headers | awk '{print $1}' | xargs basename)
    log "开始检查进程 $pid ($process_name) 的内存泄漏情况，将进行 $check_count 次检查"
    
    # 存储每次检查的内存使用情况
    local memory_usage=()
    
    # 多次检查内存使用情况
    for ((i=0; i<check_count; i++)); do
        local mem=$(ps -p $pid -o %mem --no-headers)
        local rss=$(ps -p $pid -o rss --no-headers)  #  Resident Set Size，实际使用物理内存
        memory_usage+=("$mem $rss")
        log "第 $((i+1)) 次检查: 内存使用率 $mem%，物理内存 $rss KB"
        sleep $interval
    done
    
    # 分析内存使用趋势
    local increasing=0
    local prev_rss=0
    
    for entry in "${memory_usage[@]}"; do
        local current_rss=$(echo $entry | awk '{print $2}')
        
        if [ $prev_rss -ne 0 ]; then
            # 检查内存是否持续增长(允许5%的波动)
            if (( current_rss > prev_rss * 1.05 )); then
                increasing=$((increasing + 1))
            fi
        fi
        
        prev_rss=$current_rss
    done
    
    # 如果大部分检查显示内存增长，则判断为可能存在内存泄漏
    if (( increasing >= check_count * 0.6 )); then
        log "进程 $pid ($process_name) 检测到内存泄漏迹象"
        return 0  # 存在内存泄漏
    else
        log "进程 $pid ($process_name) 未检测到明显的内存泄漏"
        return 1  # 不存在内存泄漏
    fi
}

# 处理内存泄漏进程（增强Java进程处理）
handle_memory_leak_process() {
    local pid=$1
    local user=$2
    local mem=$3
    local cmd=$4
    
    log "开始处理可能存在内存泄漏的进程: PID=$pid, USER=$user, 内存使用率=$mem%, COMMAND=$cmd"
    
    # 提取进程名
    local process_name=$(echo "$cmd" | awk '{print $1}' | xargs basename)
    
    # 特殊处理Java进程 - 创建堆转储用于分析
    if echo "$cmd" | grep -q "java"; then
        create_java_heap_dump $pid $process_name
    fi
    
    # 查找对应的重启命令
    local restart_cmd=""
    for entry in "${SERVICE_RESTART_CMD[@]}"; do
        if echo "$entry" | grep -q "^$process_name:"; then
            restart_cmd=$(echo "$entry" | cut -d':' -f2-)
            break
        fi
    done
    
    # Java进程特殊匹配（如果未找到精确匹配）
    if [ -z "$restart_cmd" ] && echo "$cmd" | grep -q "java"; then
        for entry in "${SERVICE_RESTART_CMD[@]}"; do
            if echo "$entry" | grep -q "^java:"; then
                restart_cmd=$(echo "$entry" | cut -d':' -f2-)
                break
            fi
        done
    fi
    
    # 如果有定义重启命令，使用该命令重启服务
    if [ -n "$restart_cmd" ]; then
        log "使用预定义命令重启服务 $process_name: $restart_cmd"
        eval $restart_cmd
        if [ $? -eq 0 ]; then
            log "服务 $process_name 重启成功"
            return 0
        else
            log "服务 $process_name 重启失败"
        fi
    fi
    
    # 如果没有预定义的重启命令，尝试优雅终止进程
    log "尝试优雅终止进程 $pid"
    kill -TERM $pid 2>/dev/null
    sleep 5
    
    # 检查进程是否已终止
    if ps -p $pid > /dev/null; then
        # 如果仍然运行，发送SIGKILL强制终止
        log "进程 $pid 未能优雅终止，尝试强制终止"
        kill -KILL $pid 2>/dev/null
        sleep 2
        
        if ps -p $pid > /dev/null; then
            log "无法终止进程 $pid，需要人工干预"
            return 1
        else
            log "进程 $pid 已强制终止"
        fi
    else
        log "进程 $pid 已优雅终止"
    fi
    
    # 检查进程是否会自动重启
    sleep 10
    if ! ps -ef | grep -v grep | grep -q "$cmd"; then
        log "警告: 进程 $pid 终止后未自动重启，可能需要人工干预启动"
        return 1
    else
        log "进程 $cmd 已自动重启"
        return 0
    fi
}

# 主函数
main() {
    log "内存泄漏修复脚本启动"
    log "配置: 系统内存阈值=$MEMORY_THRESHOLD%, 进程内存阈值=$PROCESS_MEM_THRESHOLD%"
    
    # 初始化Java堆转储目录
    init_java_heap_dump_dir
    
    local system_memory=$(get_system_memory_usage)
    log "当前系统内存使用率: $system_memory%"
    
    # 检查系统内存是否超过阈值
    if (( $(echo "$system_memory > $MEMORY_THRESHOLD" | bc -l) )); then
        log "系统内存使用率超过阈值 $MEMORY_THRESHOLD%，开始检查高内存占用进程"
        
        local processes=$(get_high_memory_processes)
        
        if [ -z "$processes" ]; then
            log "未发现内存使用率超过 $PROCESS_MEM_THRESHOLD% 的进程"
            return 0
        fi
        
        # 处理每个高内存占用进程
        echo "$processes" | while read -r pid user mem rss vsize cmd; do
            # 检查进程是否存在内存泄漏
            if check_memory_leak $pid $LEAK_CHECK_COUNT $((CHECK_INTERVAL / LEAK_CHECK_COUNT)); then
                # 存在内存泄漏，进行处理
                handle_memory_leak_process $pid $user $mem "$cmd"
            fi
        done
        
        # 记录处理后的系统内存使用率
        local final_memory=$(get_system_memory_usage)
        log "处理完成，当前系统内存使用率: $final_memory%"
        echo "初始内存使用率: $system_memory%"
        echo "处理后内存使用率: $final_memory%"
    else
        log "系统内存使用率在正常范围内 ($system_memory% <= $MEMORY_THRESHOLD%)，无需处理"
        echo "当前系统内存使用率: $system_memory%"
    fi
}

# 启动主函数
main
